package com.example.performance.r2dbc.compare;

import com.example.performance.r2dbc.entity.JDBCRowMapper;
import com.example.performance.r2dbc.entity.RecordDb;
import com.example.performance.r2dbc.utils.PerformanceCollector;
import com.example.performance.r2dbc.utils.TestSuite;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import org.springframework.util.StopWatch;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @author zpMeng
 * @since 6/19/2025
 */
@Slf4j
@Service
public class JdbcService {

    @Resource
    private JdbcTemplate jdbcTemplate;

    public static final DataSource dataSource = DataSourceBuilder
            .create()
            .url("jdbc:postgresql://10.86.0.7:5432/fmuser_atlas_demo")
            .username("fmuser_atlas_demo")
            .password("fm")
            .build();
    public static final JdbcTemplate JDBC_TEMPLATE = new JdbcTemplate(dataSource);

    public void testSerialPerformance() {

        Map<Long, Integer> nodeSizeWithTestTimes = TestSuite.prepareSerialTestSuite();
        System.out.println(nodeSizeWithTestTimes);

        for (Map.Entry<Long, Integer> nodeSizeWithTestTime : nodeSizeWithTestTimes.entrySet()) {
            Long nodeSize = nodeSizeWithTestTime.getKey();
            Integer testTime = nodeSizeWithTestTime.getValue();

            ArrayList<Double> durations = new ArrayList<>();
            for (int i = 0; i < testTime; i++) {
                double duration = executeOneJdbcQuery(0, nodeSize);
                durations.add(duration);
            }

            Double avg = PerformanceCollector.getAverage(durations);

            System.out.printf("nodeSize: %d -- avg time: %f %n", nodeSize, avg);
        }

    }

    public void testParallelPerformance(int start, int end, int nodeSize) throws InterruptedException {
        Map<Integer, Integer> concurrencyWithTestTimes = TestSuite.prepareParallelTestSuite(start, end, nodeSize);
        System.out.println(concurrencyWithTestTimes);

        for (Map.Entry<Integer, Integer> entry : concurrencyWithTestTimes.entrySet()) {
            Integer concurrency = entry.getKey();
            Integer testTime = entry.getValue();

            testOneConcurrencyPerformance(testTime, concurrency, nodeSize);
        }
    }

    public void testOneConcurrencyPerformance(int count, int concurrency, int nodeSize) throws InterruptedException {
        ArrayList<Double> totalDurations = new ArrayList<>();

        for (int i = 1; i <= count; i++) {
            try (ExecutorService executor = Executors.newFixedThreadPool(concurrency);) {
                List<Double> singleTestTaskTimes = executeOneParallelTest(executor, concurrency, nodeSize);
                totalDurations.addAll(singleTestTaskTimes);
            }
        }

        PerformanceCollector.printParallel(totalDurations, concurrency);
    }

    public List<Double> executeOneParallelTest(ExecutorService executorService, int concurrency, int nodeSize) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(concurrency);
        List<Future<Double>> futures = new ArrayList<>(concurrency);
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();

        for (int i = 0; i < concurrency; i++) {
            int version = concurrency % 20;

            Future<Double> currentTask = executorService.submit(() -> {
                try {
                    return executeOneJdbcQuery(version, (long) nodeSize);
                } finally {
                    latch.countDown();
                }
            });
            futures.add(currentTask);
        }

        latch.await();
        stopWatch.stop();

        // 收集每个任务的执行时间
        List<Double> taskTimes = new ArrayList<>();
        for (Future<Double> future : futures) {
            try {
                taskTimes.add(future.get());
            } catch (ExecutionException e) {
                System.err.println("Task failed: " + e.getCause());
            }
        }

        return taskTimes;
    }

    public double executeOneJdbcQuery(int version, Long nodeSize) {
        String sql = "select * from DS_AGG_427_%d where node_id <= %d".formatted(version, nodeSize);

        StopWatch stopWatch = new StopWatch();
        stopWatch.start("JDBC query");

        List<RecordDb> recordList = jdbcTemplate.query(sql, new JDBCRowMapper());
        // System.out.println(recordList.size());

        stopWatch.stop();
        return stopWatch.getTotalTimeSeconds();
    }
}
